from export.export import export_to_excel
from httpx.httpx import HTTPX
import pandas as pd
from config.settings import AWVS_URL, AWVS_API, PACKAGE_PATH
from awvs.awvs import AWVS
from cyber_search_engine.shodan import Shodan
from cyber_search_engine.zoomeye import Zoomeye
from rapiddns.rapiddns import RapidDNS
from goby.goby import Goby
from utils.excel import get_ip_from_excel, get_subdomains_from_excel
from get_sheet import get_subdomains

def add_awvs_targets(save_path):
    awvs = AWVS()
    db = pd.read_excel(save_path, sheet_name='存活域名')
    alive_domains = list(db['url'])
    awvs.run(alive_domains, save_path=save_path)

def extract_vulns():
    awvs = AWVS()
    awvs.awvs.workflow_vuln_details()
    
def start_awvs_targets(save_path):
    db = pd.read_excel(save_path, sheet_name='AWVS(target)')
    target_ids = list(db['target_id'])
    awvs = AWVS()
    for target_id in target_ids:
        print(awvs.awvs.start_scan(target_id))

def complete_awvs_targets(save_path):
    db = pd.read_excel(save_path, sheet_name='AWVS(target)')
    target_ids = list(db['target_id'])
    awvs = AWVS()
    awvs.awvs.workflow_vuln_details(target_ids, save_path=save_path)

def resume_ip_infos(file_path, sheet_name):
    db = pd.read_excel(file_path, sheet_name=sheet_name)
    ips = db.loc[db['解析类型']== 'A']
    ips = list(ips['解析结果'])
    goby = Goby()
    goby.run(ips, save_path=file_path)
    zoom = Zoomeye(ips)
    zoom.run(file_path)
    shodan = Shodan(ips)
    shodan.run(file_path)

def complete_ip_infos(save_path):
    ips = get_ip_from_excel(save_path)
    zoom = Zoomeye(ips)
    zoom.run(save_path)
    
def complete_goby_infos(save_path):
    goby = Goby()
    db = pd.read_excel(save_path, sheet_name='goby')
    task_ids = list(db['goby_task'])
    goby.export_assets(task_ids[0], save_path)

def complete_httpx_infos(save_path):
    domains = get_subdomains_from_excel(save_path)
    httpx_exec_path = PACKAGE_PATH / 'thirdparty/httpx'
    httpx = HTTPX(httpx_exec_path, domains)
    httpx.run(save_path)
    
def complete_subdomain_infos(save_path):
    db = pd.read_excel(save_path, sheet_name='ICP备案信息')
    domains = list(db['域名'])
    rapiddns = RapidDNS(domains)
    rapiddns.run(save_path)

def complete_subdomain_all_infos(save_path):
    db = pd.read_excel(save_path, sheet_name='ICP备案信息')
    domains = list(db['域名'])
    for domain in domains:
        try:
            result = get_subdomains(domain)
            print(f'-[]{domain} 子域名数量: {len(result)}')
            export_to_excel(result, xlsx_name=domain, sheet_name='domain-' + domain, save_path=save_path)
        except Exception as e:
            print(e)
            continue

if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='Domain or File')
    parser.add_argument('--domain', type=str, help='domain')
    parser.add_argument('--file', type=str, help='file path')
    parser.add_argument('--complete-awvs', dest='awvs', action='store_true')
    parser.add_argument('--complete-goby', dest='goby', action='store_true')
    parser.add_argument('--complete-subdomains', dest='subdomains', action='store_true')
    parser.add_argument('--complete-subdomains-all', dest='subdomains_all', action='store_true')
    parser.add_argument('--complete-httpx', dest='httpx', action='store_true')
    parser.add_argument('--complete-zoomeye', dest='zoomeye', action='store_true')
    parser.add_argument('--start-awvs', dest='start_awvs', action='store_true')
    parser.add_argument('--resume-ip', dest='resume_ip', action='store_true')
    args = parser.parse_args()

    if args.domain:
        save_path = f'res/{args.domain}.xlsx'
    
        if args.goby:    
            complete_goby_infos(save_path)
        
        if args.awvs:
            complete_awvs_targets(save_path)
            
        if args.start_awvs:
            start_awvs_targets(save_path)
        
        if args.resume_ip:
            resume_ip_infos(save_path, args.domain)

        if args.subdomains:
            complete_subdomain_infos(save_path)
        
        if args.subdomains_all:
            complete_subdomain_all_infos(save_path)
            
        if args.httpx:
            complete_httpx_infos(save_path)
            
        if args.zoomeye:
            complete_ip_infos(save_path)
